import struct
import json
from exception import InvalidPackageLengthError

defaultVer = 1
defaultServer = 0
defaultMagic = 0
packHeadLength = 8
magicMap = {
    "json": 1,
    "bytes": 2,
    "disconnect": 3
}
headFormat = """PackageHead:
version: {}
magic:   {}
server:  {}
len:     {}"""
jsonPackFormat = """JsonPackage:
version: {}
magic:   {}
server:  {}
len:     {}
json:
{}"""
bytePackFormat = """BytePackage:
version: {}
magic:   {}
server:  {}
len:     {}
head:
{}
data:
{}"""
encoding = "UTF-8"


class PackageHead():
    """
    应用层包头部处理

    格式：
    1B ver 版本
    1B magic
    2B server 服务器功能
    4B len 数据包长度
    """

    def __init__(self, version: int=defaultVer, magic: int=defaultMagic, server: int=defaultServer, length: int=0):
        self._setHeadData(version, magic, server, length)

    def readBinary(self, data: bytes):
        self._readHeadBinary(data)

    def _readHeadBinary(self, data: bytes):
        """
        data 长度可能超过 packHeadLength
        """
        length = len(data)
        if length < packHeadLength:
            raise InvalidPackageLengthError(length, packHeadLength, PackageHead)
        if None == data:
            self.ver = None
            self.magic = None
            self.server = None
            self.len = None
            return
        (self.ver, self.magic, self.server, self.len) = struct.unpack('>BBHI', data[:packHeadLength])

    def _setHeadData(self, version: int, magic: int, server: int, length: int):
        self.ver = version
        self.magic = magic
        self.server = server
        self.len = length

    def setHead(self, version: int, magic: int, server: int):
        self.ver = version
        self.magic = magic
        self.server = server

    def __str__(self):
        """
        转换为 str，用于打印输出
        """
        return headFormat.format(self.ver, self.magic, self.server, self.len)

    def __bytes__(self):
        """
        转换为字节串 bytes
        """
        return struct.pack('>BBHI', self.ver, self.magic, self.server, self.len)

    def __eq__(self, other: 'PackageHead') -> bool:
        """
        判断相等
        """
        if self.ver == other.ver and self.magic == other.magic \
                and self.server == other.server and self.len == other.len:
            return True
        else:
            return False

class JsonPackage(PackageHead):
    """
    JSON 应用层包
    """

    def __init__(self, data=None):
        """
        初始化 JsonBody
        data 不是二进制数据，而是要转换为二进制的数据
        """
        self._setHeadData(defaultVer, magicMap['json'], defaultServer, 0)
        self.data = data

    def readBinary(self, data: bytes):
        self._readHeadBinary(data)
        length = len(data)
        if length != self.len:
            raise InvalidPackageLengthError(length, self.len, JsonPackage)
        self.data = json.loads(data[packHeadLength:].decode(encoding))

    def setData(self, data):
        self.data = data

    def __str__(self):
        return jsonPackFormat.format(
                self.ver,
                self.magic,
                self.server,
                self.len,
                json.dumps(self.data)
            )

    def __bytes__(self):
        binData = json.dumps(self.data).encode(encoding)
        self.len = len(binData) + packHeadLength
        return b''.join([super(JsonPackage, self).__bytes__(), binData])

    def __eq__(self, other: 'JsonPackage') -> bool:
        if super(JsonPackage, self).__eq__(other) \
                and self.data == other.data:
            return True
        else:
            return False

class BytePackage(PackageHead):
    """
    二进制应用层包
    
    包体包括 JSON 字段字节长度、JSON 头和二进制数据 三部份
    """

    def __init__(self, jsonHead: dict=None, data: bytes=None):
        self._setHeadData(defaultVer, magicMap['bytes'], defaultServer, 0)
        self.setData(jsonHead, data)

    def readBinary(self, data: bytes):
        self._readHeadBinary(data)
        length = len(data)
        if length != self.len:
            raise InvalidPackageLengthError(length, self.len, BytePackage)
        cur = packHeadLength
        (jsonLen, ) = struct.unpack('>I', data[cur:cur + 4])
        cur += 4
        self.head = json.loads(data[cur:cur + jsonLen].decode(encoding))
        cur += jsonLen
        self.data = data[cur:][:]

    def setData(self, jsonHead: dict, data: bytes):
        self.head = jsonHead
        if None == data:
            self.data = None
        else:
            self.data = data[:]

    def __str__(self):
        return bytePackFormat.format(
                self.ver,
                self.magic,
                self.server,
                self.len,
                json.dumps(self.head),
                self.data.hex()
            )

    def __bytes__(self):
        """
        JSON 字段字节长度：4 字节
        JSON 字段
        二进制数据
        """
        jsonBytes = json.dumps(self.head).encode(encoding)
        jsonLen = len(jsonBytes)
        lenBytes = struct.pack('>I', jsonLen)
        self.len = packHeadLength + 4 + jsonLen + len(self.data)
        headBytes = super(BytePackage, self).__bytes__()
        return b''.join([headBytes, lenBytes, jsonBytes, self.data])

    def __eq__(self, other: 'BytePackage') -> bool:
        if super(BytePackage, self).__eq__(other) \
                and self.data == other.data \
                and self.head == other.head:
            return True
        else:
            return False

magicRevMap = {
        1: JsonPackage,
        2: BytePackage,
        3: PackageHead
}

def readBinary(data: bytes):
    """
    data 必须正好是一个包的数据，否则抛出错误

    return:
    (JsonPackage or BytePackage, bytes)
    """
    head = PackageHead()
    head._readHeadBinary(data)
    length = len(data)
    if length != head.len:
        raise InvalidPackageLengthError(length, head.len, PackageHead)
    if head.magic in magicRevMap:
        result = magicRevMap[head.magic]()
        result.readBinary(data[:head.len])
        return result
    else:
        raise Exception("Invalid magic number")
